import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder

from assignment_4.src.scripts import getinfo_d, get_min_d_index, create_new_data
from data_structures.info import Info


#
# # 创建DataFrame
# data = {
#     'age': ['youth', 'youth', 'middle_aged', 'senior', 'senior', 'senior', 'middle_aged', 'youth', 'youth', 'senior',
#             'youth', 'middle_aged', 'middle_aged', 'senior'],
#     'income': ['high', 'high', 'high', 'medium', 'low', 'low', 'low', 'medium', 'low', 'medium', 'medium', 'medium',
#                'high', 'medium'],
#     'student': ['no', 'no', 'no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no'],
#     'credit_rating': ['fair', 'excellent', 'fair', 'fair', 'fair', 'excellent', 'excellent', 'fair', 'fair', 'fair',
#                       'excellent', 'excellent', 'fair', 'excellent'],
#     'buys_computer': ['no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no', 'yes', 'yes', 'yes', 'yes', 'yes', 'no']
# }
#
# df = pd.DataFrame(data)
#
# # 使用LabelEncoder将类别数据转化为数值
# label_encoder = LabelEncoder()
# df_encoded = df.apply(label_encoder.fit_transform)
#
# # 特征与目标变量
# X = df_encoded.drop('buys_computer', axis=1)
# y = df_encoded['buys_computer']
#
# # 数据分割
# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
#
# # 创建并训练决策树
# clf = DecisionTreeClassifier(random_state=42)
# clf.fit(X_train, y_train)
#
# # 评估模型
# accuracy = clf.score(X_test, y_test)
# print(f"Accuracy: {accuracy:.2f}")


def create_tree(T):
    buys_computer_d = getinfo_d(T, 'buys_computer')
    if buys_computer_d == 0:
        print('此分支无需再划分')
        return 1
    else:
        table_head = T.columns.tolist()  # 获取表头
        min_d_index = get_min_d_index(table_head, T, 'buys_computer')

        index = T[table_head[min_d_index]].tolist()
        index = list(set(index))

        for a in index:  # 进行分割
            new_head = [item for item in table_head if item != table_head[min_d_index]]
            new_lines = []

            for b in T.values:
                if a in b:
                    temp = 0
                    temp_list = []
                    for c in b:
                        if temp != min_d_index:
                            temp_list.append(c)
                        temp += 1
                    new_lines.append(temp_list)

            new_t = create_new_data(new_head, new_lines)
            print('\n')
            print(f'以{table_head[min_d_index]}-{a}进行划分')
            print(new_t)
            create_tree(new_t)
        return 1


if __name__ == '__main__':
    data = {
        'age': ['youth', 'youth', 'middle_aged', 'senior', 'senior', 'senior', 'middle_aged', 'youth', 'youth',
                'senior',
                'youth', 'middle_aged', 'middle_aged', 'senior'],
        'income': ['high', 'high', 'high', 'medium', 'low', 'low', 'low', 'medium', 'low', 'medium', 'medium', 'medium',
                   'high', 'medium'],
        'student': ['no', 'no', 'no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no'],
        'credit_rating': ['fair', 'excellent', 'fair', 'fair', 'fair', 'excellent', 'excellent', 'fair', 'fair', 'fair',
                          'excellent', 'excellent', 'fair', 'excellent'],
        'buys_computer': ['no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no', 'yes', 'yes', 'yes', 'yes', 'yes', 'no']
    }

    T = pd.DataFrame(data)
    print(T)
    create_tree(T)
